#!/usr/bin/env python3
import rospy
import threading
import numpy as np
from std_msgs.msg import UInt32MultiArray
from sensor_msgs.msg import Image
from event_recorder.msg import CameraData
# 是否为debug模式
debugMode = True

# 定义发布者
publisherInitialized = False
nodePub = None

# 定义订阅者
subscriberInitialized = False
nodeSub = None

eventBuffer = None

from cv_bridge import CvBridge
br = CvBridge()

"""
该节点的回调函数
@params:
    msg: EventVecInline 收到的（多个事件）消息
"""
def nodeSubCallback(msg: CameraData) -> None:
    global publisherInitialized, nodePub, eventBuffer
    eventData = np.array(msg.events, dtype = "uint32")
    imuData = np.array(msg.imus, dtype = "float64")
    if eventData.shape[0]:
        eventData = eventData.reshape((eventData.shape[0] // 5, 5)) # [t_sec, t_nsec, p, y, x]
        if eventBuffer is None:
            eventBuffer = eventData
        else:
            eventBuffer = np.vstack([eventData, eventBuffer])
    if imuData.shape[0]:
        imuData = imuData.reshape((imuData.shape[0] // 4, 4)) # [t, wz, wy, wx]
    

"""
订阅者的线程函数，用于执行订阅者线程
"""
def subThreadFunc() -> None:
    global debugMode, subscriberInitialized, nodeSub
    nodeSub = rospy.Subscriber("/event_recorder/sub_event_res", CameraData, nodeSubCallback)  # 订阅事件消息
    subscriberInitialized = True
    rospy.spin()  # 开始循环监听消息


"""
发布者的线程函数，用于发布本结点的执行结果
"""
def pubThreadFunc() -> None:
    global debugMode, publisherInitialized, nodePub, imgPub, eventBuffer
    nodePub = rospy.Publisher("/event_recorder/event_res_img", Image, queue_size = 1)
    publisherInitialized = True
    frameRate = rospy.Rate(30)
    while not rospy.is_shutdown():
        res = np.zeros((800, 1280, 3), dtype = "uint8")
        if eventBuffer is None:
            nodePub.publish(br.cv2_to_imgmsg(res, encoding = "bgr8"))
            frameRate.sleep()
            continue
        savedEvents = eventBuffer.copy()
        eventBuffer = None
        savedEvents = savedEvents[::10]
        pos = savedEvents[savedEvents[:, 2] == 1, 3:5]
        neg = savedEvents[savedEvents[:, 2] == 0, 3:5]
        if pos.shape[0]:
            res[pos[:, 0], pos[:, 1], 2] = 255
        if neg.shape[0]:
            res[neg[:, 0], neg[:, 1], 0] = 255
        nodePub.publish(br.cv2_to_imgmsg(res, encoding = "bgr8"))
        frameRate.sleep()


if __name__ == "__main__":
    rospy.init_node("view", anonymous=True)  # 初始化ROS环境

    subT = threading.Thread(target=subThreadFunc)  # 定义订阅者线程
    pubT = threading.Thread(target=pubThreadFunc)  # 定义发布者线程

    subT.start()  # 启动订阅者线程
    pubT.start()  # 启动发布者线程

    subT.join()  # 加入订阅者线程
    pubT.join()  # 加入发布者线程
